import os
from gradio.layouts import row
import pandas as pd
import gradio as gr

import plotly.express as px
from plotly.graph_objs import Figure

from utils import EpisodeItem, TransactionFlowAgentItem



class ShowData:
    def __init__(self):
        self.path = "../"
        # epitem 为交易细项
        self.ep_item = EpisodeItem()
        self.tf_item = TransactionFlowAgentItem() # tf item 为交易大项目（资金流量项目）
        # 资产负债表地址和资金流量表地址
        self.path_balance_sheet = [f"{self.path}{s}" for s in os.listdir(self.path) if '资产负债表' in s and 'xlsx' in s]
        self.path_transaction_flow = [f"{self.path}{s}" for s in os.listdir(self.path) if '资金流量表' in s and 'xlsx' in s]
        # 资产负债表和资金流量表数据，以字典形式存储，键为年份，值为数据表
        print(self.path_balance_sheet)
        self.bs_data = pd.read_excel(self.path_balance_sheet[0], sheet_name=None, index_col=0)
        self.tf_data = pd.read_excel(self.path_transaction_flow[0], sheet_name=None, index_col=0)
        # 把数据表的行索引名称作为第一列，这样在gradio中显示时，第一列就是行索引名称，方便选择
        self.bs_data_ = {k: self.transfer_data(v) for k, v in self.bs_data.items()}
        self.tf_data_ = {k: self.transfer_data(v) for k, v in self.tf_data.items()}
        # bs_periods, tf_periods 为资产负债表和资金流量表的年份，periods 为两者的并集
        self.bs_periods = list(self.bs_data.keys())
        self.tf_periods = list(self.tf_data.keys())
        self.periods = self.bs_periods
        # bs_idx, bs_col 为资产负债表的行索引和列索引，tf_idx, tf_col 为资金流量表的行索引和列索引
        self.bs_idx = self.bs_data_[self.periods[0]].index.tolist()
        self.bs_col = self.bs_data_[self.periods[0]].columns.tolist()
        self.tf_idx = self.tf_data_[self.periods[0]].index.tolist()
        self.tf_col = self.tf_data_[self.periods[0]].columns.tolist()
        self.int_periods = [int(s) for s in self.periods]
    
    @staticmethod    
    def transfer_data(df:pd.DataFrame):
        """
        把数据表的行索引名称作为第一列
        """
        df.insert(0, 'Index', df.index)  # 添加索引列，将行索引名称作为第一列
        return df

    def update_bs_tf_tables(self, year_ratio):
        """
        根据所选的年份更新资产负债表和资金流量表。
        """
        selected_year = str(year_ratio)
        bs_data = self.bs_data_[selected_year].round(1)
        tf_data = self.tf_data_[selected_year].round(1)
        return bs_data, tf_data

    def change_tf_agent_cols(self, tf_col_select):
        """
        选择资金流量列名称后，匹配有数值的流量项目
        """
        print(f">>> 你选择的是：（资金流量表） {tf_col_select}")
        df_tf_last_year = self.tf_data_[self.periods[-1]]
        col_se_tf = df_tf_last_year.iloc[:,1:].loc[:,tf_col_select]
        col_se_tf_float = {k:float(v) for k,v in col_se_tf.items()}
        idx_tf = [k for k,v in col_se_tf_float.items() if v!= 0]
        print(f"资金流量表列:{str(col_se_tf)}")
        return gr.CheckboxGroup(choices=self.tf_idx, 
                                value=idx_tf, 
                                interactive=True)

    def change_bs_tf_col(self, bs_tf_col_select):
        """
        选择资产名称后，根据资产名称选择部门
        """
        print(f">>> 你选择的是：（资产负债表） {bs_tf_col_select}")
        df_bs_last_year = self.bs_data_[self.periods[-1]]
        col_se_bs = df_bs_last_year.iloc[:,1:].loc[:,bs_tf_col_select]
        col_se_bs_float = {k:float(v) for k,v in col_se_bs.items()}
        # 根据bs_tf_col_select 选择最后一年的数据, 过滤掉值为0的列
        cols_bs = [k for k,v in col_se_bs_float.items() if v!= 0]
        print(f"资产负债表列:{str(cols_bs)}")
        return gr.CheckboxGroup(choices=col_se_bs.index.to_list(), 
                                value=cols_bs, 
                                interactive=True)
        
    def change_bs_item(self, bs_row_item_select):
        """
        选择资产名称后，根据资产名称选择部门
        """
        print(f">>> 你选择的是：（资产负债表） {bs_row_item_select}")
        df_bs_first_year = self.bs_data_[self.periods[0]]
        row_se = df_bs_first_year.iloc[:,1:].loc[bs_row_item_select]
        print(f"第一年数据:", dict(row_se))
        row_se_float = {k:float(v) for k,v in row_se.items()}
        cols = [k for k,v in row_se_float.items() if v != 0]
        print(cols)
        return gr.CheckboxGroup(choices=row_se.index.to_list(), 
                                value=cols, interactive=True)
        
    def change_tf_item(self, tf_row_item_select):
        """
        选择资金流量项目后，根据资金流量项目选择部门
        """
        print(f">>> 你选择的是:资金（交易）流量表 {tf_row_item_select}")
        # 这里如果索引periods[0]reserve，cash advance可能还没发生
        df_bs_first_year = self.tf_data_[self.periods[-1]]
        row_se = df_bs_first_year.iloc[:,1:].loc[tf_row_item_select]
        print(f"第一年数据:", dict(row_se))
        row_se_float = {k:float(v) for k,v in row_se.items()}
        cols = [k for k,v in row_se_float.items() if v != 0]
        print(cols)
        return gr.CheckboxGroup(choices=row_se.index.to_list(), 
                                value=cols, interactive=True)
    
    def plot_by_bs_item(self, bs_row_item_select, bs_col_item_select):
        """
        选择资产名称和部门后，绘制时序图
        """
        dic = {}
        for bs_col_select in bs_col_item_select:
            print(f">>> 你选择的是 {bs_row_item_select} 和 {bs_col_select}")
            se = {k:v.loc[bs_row_item_select, bs_col_select] for k,v in self.bs_data_.items()}
            dic[bs_col_select] = se
        new_df = pd.DataFrame(dic)
        print(new_df)
        fig = px.line(new_df, x=new_df.index, y=new_df.columns)
        return fig
    
    def plot_by_bs_item_multi_row(self, bs_row_item_select, bs_col_item_select):
        """
        选择资产名称和部门后，绘制时序图
        """
        dic = {}
        for row_col_select in bs_row_item_select:
            print(f">>> 你选择的是 {row_col_select} 和 {bs_col_item_select}")
            se = {k:v.loc[row_col_select, bs_col_item_select] for k,v in self.bs_data_.items()}
            dic[row_col_select] = se
        new_df = pd.DataFrame(dic)
        print(new_df)
        fig = px.line(new_df, x=new_df.index, y=new_df.columns)
        return fig
    
    def plot_by_tf_item_multi_row(self, tf_row_item_select, tf_col_item_select):
        """
        选择流量名称和部门后，绘制时序图
        """
        dic = {}
        for row_col_select in tf_row_item_select: # 这里为列表
            print(f">>> 你选择的是 {row_col_select} 和 {tf_col_item_select}")
            se = {k:v.loc[row_col_select, tf_col_item_select] for k,v in self.tf_data_.items()}
            dic[row_col_select] = se
        new_df = pd.DataFrame(dic)
        print(new_df)
        fig = px.line(new_df, x=new_df.index, y=new_df.columns)
        return fig
    
    def plot_by_tf_item(self, tf_row_item_select, tf_col_item_select):
        """
        选择资金流量项目和部门后，绘制时序图
        """
        dic = {}
        for tf_col_select in tf_col_item_select:
            print(f">>> 你选择的是 {tf_row_item_select} 和 {tf_col_select}")
            se = {k:v.loc[tf_row_item_select, tf_col_select] for k,v in self.tf_data_.items()}
            dic[tf_col_select] = se
        new_df = pd.DataFrame(dic)
        print(new_df)
        fig = px.line(new_df, x=new_df.index, y=new_df.columns)
        return fig
        
    
    @staticmethod
    def plot_se(se):
        se = pd.Series(se)
        fig = px.bar(se, x=se.index, y=se.values)
        return fig
    
    def plot_bs(self, bs_row_select, bs_col_select):
        se = {k:v.loc[bs_row_select, bs_col_select] for k,v in self.bs_data_.items()}
        fig = self.plot_se(se)
        return fig
        
    def html(self):
        with gr.Blocks(css=".dataframe tr { height: 30px; }") as demo:  # 添加 CSS 样式设置行高
            gr.Markdown("# Agent Based - Stock Flow Consistent Model(多主体存量流量一致模型)")
            with gr.Tab("数据全览"):
                gr.Markdown("## 选择年份，自动更新资产负债表和资金流量表")
                year_ratio = gr.Slider(
                    minimum=min(self.int_periods),
                    maximum=max(self.int_periods),
                    step=1,
                    value=min(self.int_periods),
                    label="选择年份",
                    interactive=True,
                    show_label=True,
                )
                gr.Markdown("## 资产负债数据")
                # 获取数据的索引作为行名称
                bs_row_names = self.bs_data_[self.periods[0]].index.tolist()
                bs_col_names = self.bs_data_[self.periods[0]].columns.tolist()
                bs_table = gr.Dataframe(
                    value=self.bs_data_[self.periods[0]].round(1),
                    datatype="number",
                    interactive=True,
                    col_count=(len(bs_col_names), "fixed"),
                    row_count=(len(bs_row_names), "fixed"),  # 设置行数 
                )
                gr.Markdown("## 资金流量数据")
                # 获取数据的索引作为行名称
                tf_row_names = self.tf_data_[self.periods[0]].index.tolist()
                tf_col_names = self.tf_data_[self.periods[0]].columns.tolist()
                tf_table = gr.Dataframe(
                    value=self.tf_data_[self.periods[0]].round(1),
                    datatype="number",
                    interactive=True,
                    col_count=(len(tf_col_names), "fixed"),
                    row_count=(len(tf_row_names), "fixed"),  # 设置行数 
                    # 移除不支持的 height 参数
                )
            with gr.Tab("按资产类别"):
                gr.Markdown("## 选择资产类别，自动输出有数据的部门")
                # 获取数据的索引作为行名称
                with gr.Row():
                    bs_row_item_select = gr.Radio(
                            choices=self.bs_idx,
                            value=self.bs_idx[0],
                            label="选择资产名称",
                            interactive=True,
                            show_label=True,
                            scale=2,
                        )
                    bs_col_item_select = gr.CheckboxGroup(
                            choices=self.bs_idx,
                            value=self.bs_idx[0],
                            label="选择部门名称（自动）",
                            interactive=True,
                            show_label=True,
                            scale=2,
                        )
                plot_by_bs_item_bt = gr.Button("查看时序数据")
                plot_by_bs_item = gr.Plot()
            with gr.Tab("按流量类别"):
                gr.Markdown("## 选择流量类别，自动输出有数据的部门")
                tf_row_item_select = gr.Radio(
                        choices=self.tf_idx,
                        value=self.tf_idx[0],
                        label="选择流量名称",
                        interactive=True,
                        show_label=True,
                        scale=2,
                    )
                tf_col_item_select = gr.CheckboxGroup(
                        choices=self.tf_col,
                        value=self.tf_col[0],
                        label="选择部门名称（自动）",
                        interactive=True,
                        show_label=True,
                        scale=2,
                    )
                plot_by_tf_item_bt = gr.Button("查看时序数据")
                plot_by_tf_item = gr.Plot()
            
            with gr.Tab("按部门看资产负债"):
                gr.Markdown("## 选择部门，自动输出该部门所有资产负债项目")
                bs_tf_col_item_select = gr.Radio(
                        choices=self.bs_col[1:],
                        value=self.bs_col[1:][0],
                        label="选择部门名称",
                        interactive=True,
                        show_label=True,
                        scale=2,
                    )
                with gr.Row():
                    bs_row_answer_by_col = gr.CheckboxGroup(
                            choices=self.bs_idx,
                            value=self.bs_idx[0],
                            label="选择资产名称（自动）",
                            interactive=True,
                            show_label=True,
                            scale=2,
                        )
                plot_bs_agt = gr.Button("查看时序数据")
                bs_col_item_plot = gr.Plot()
            
            with gr.Tab("按部门看资金流量"):
                gr.Markdown("## 选择部门，自动输出该部门所有资金流量项目")
                look_tf_by_agent_cols = gr.Radio(
                        choices=self.tf_col,
                        value=self.tf_col[0],
                        label="选择部门名称",
                        interactive=True,
                        show_label=True,
                        scale=2,
                    )
                answer_tf_rows_by_agent = gr.CheckboxGroup(
                            choices=self.tf_idx,
                            value=self.tf_idx[0],
                            label="选择资产名称（自动）",
                            interactive=True,
                            show_label=True,
                            scale=2,
                        )
                plot_tf_agt = gr.Button("查看时序数据")
                tf_col_item_plot = gr.Plot()
                
            look_tf_by_agent_cols.change(
                fn=self.change_tf_agent_cols,
                inputs=[look_tf_by_agent_cols],
                outputs=[answer_tf_rows_by_agent]
            )    
            
            bs_tf_col_item_select.change(
                fn=self.change_bs_tf_col,
                inputs=[bs_tf_col_item_select],
                outputs=[bs_row_answer_by_col]
            )
            
            tf_row_item_select.change(
                fn=self.change_tf_item,
                inputs=[tf_row_item_select],
                outputs=[tf_col_item_select]
            )
            # tab 按流量类别
            plot_by_tf_item_bt.click(
                fn=self.plot_by_tf_item,
                inputs=[tf_row_item_select, tf_col_item_select],
                outputs=[plot_by_tf_item]
            )
            # tab 按资产类别
            plot_by_bs_item_bt.click(
                fn=self.plot_by_bs_item,
                inputs=[bs_row_item_select, bs_col_item_select],
                outputs=[plot_by_bs_item]
            )
            # 按部门看资产负债
            plot_bs_agt.click(
                fn=self.plot_by_bs_item_multi_row,
                inputs=[bs_row_answer_by_col, bs_tf_col_item_select],
                outputs=[bs_col_item_plot]
            )
            plot_tf_agt.click(
                fn=self.plot_by_tf_item_multi_row,
                inputs=[answer_tf_rows_by_agent, look_tf_by_agent_cols],
                outputs=[tf_col_item_plot]
            )
            # 改变资产名称后，自动改变部门名称（总资产负债中有这一项数据） 
            bs_row_item_select.change(
                fn=self.change_bs_item,
                inputs=[bs_row_item_select],
                outputs=[bs_col_item_select]
            )
            # 年份数据改变后，自动更新资产负债表和资金流量表    
            year_ratio.change(
                fn=self.update_bs_tf_tables,
                inputs=[year_ratio],
                outputs=[bs_table, tf_table]
            )
        return demo

sd = ShowData()
sd.html().launch()